# -*- coding: utf-8 -*-
# @Time    : 2023/7/4 19:43
# @Author  : Pan
# @Software: PyCharm
# @Project : VisualFramework
# @FileName: __init__.py
import paddle
from networks.Image2Iamge.UNet import UNet
from networks.Image2Iamge.UNetPan_v1 import UNetPanV1
from networks.Image2Iamge.UNite import ULiteNet
from networks.Image2Iamge.ULitePan import ULitePan


def image2image(config):
    model = eval(config["type"])(config)
    if "pretrained" in config and config["pretrained"] is not None:
        model.set_state_dict(paddle.load(config["pretrained"]))
    return model
